%% 1D infinite box, FDM method
clear all;
clc;
global h e_charge hbar m eps x_max y_max ax ay x y hsize;

h=6.626e-34;
e_charge=1.609e-19;
hbar=h/(2*pi);
mass=9.11e-31;
eps=8.854e-12;

x_max=1e-9;
y_max=1e-9;
ax=1e-10; %spacing between each point
ay=1e-10;
% x=a*[1:1:npoints];
x=ax:ax:x_max;
y=ay:ay:y_max;
m=length(x);
n=length(y);
hsize = m*n;
%m=1000;


Hmat=zeros(hsize,hsize);
t0x=hbar^2/(2*mass*ax^2)/e_charge;
t0y=hbar^2/(2*mass*ay^2)/e_charge;
psi=zeros(1,hsize)';

U=eye(m,n);
U_Calc=U(:);

figure(3)
mesh(x,y,U);
title('Potential Variation');
grid on;


BLC = 0;
BRC = 1;
TLC = 2;
TRC = 3;
BEN = 4;
LEN = 5;
REN = 6;
TEN = 7;
MID = 8;



%% set the node type matrix
node_type=zeros(m,n);
for i=1:m
    for j=1:n
        if i==1
            if j==1
                node_type(i,j)=BLC;
            elseif j ==n
                node_type(i,j)=BRC;
            else
                node_type(i,j)=BEN;
            end
        elseif i==m
            if j==1
                node_type(i,j)=TLC;
            elseif j ==n
                node_type(i,j)=TRC;
            else
                node_type(i,j)=TEN;
            end
        else
            if j==1
                node_type(i,j)=LEN;
            elseif j ==n
                node_type(i,j)=REN;
            else
                node_type(i,j)=MID;
            end
        end
    end
end

figure(3)
mesh(x,y,node_type);
title('NodeType');
grid on;

ntc=node_type(:);


% Hmat=diag((2*t0).* ones(1,m)) - (t0*diag(ones(1,m-1),1)) - (t0*diag(ones(1,m-1),-1));
%% prepare the HMatrix
for i=1:hsize
    if ntc(i)==BLC
        Hmat(i,i)=t0x+t0y;
        Hmat(i,i+n)=-t0y;
        Hmat(i,i+1)=-t0x;
    elseif(ntc(i) == BEN)
        Hmat(i,i)=2*t0x+t0y;
        Hmat(i,i+n)=-t0y;
        Hmat(i,i+1)=-t0x;
        Hmat(i,i-1)=-t0x;
    elseif(ntc(i) == BRC)
        Hmat(i,i)=t0x+t0y;
        Hmat(i,i+n)=-t0y;
        Hmat(i,i-1)=-t0x;
    elseif(ntc(i) == REN)
        Hmat(i,i)=t0x+2*t0y;
        Hmat(i,i+n)=-t0y;
        Hmat(i,i-n)=-t0y;
        Hmat(i,i-1)=-t0x;
    elseif(ntc(i) == LEN)
        Hmat(i,i)=t0x+2*t0y;
        Hmat(i,i+n)=-t0y;
        Hmat(i,i-n)=-t0y;
        Hmat(i,i+1)=-t0x;
    elseif (ntc(i) == TEN)
        Hmat(i,i)=2*t0x+t0y;
        Hmat(i,i+1)=-t0x;
        Hmat(i,i-n)=-t0y;
        Hmat(i,i-1)=-t0x;
    elseif(ntc(i) == TLC)
        Hmat(i,i)=t0x+t0y;
        Hmat(i,i-n)=-t0y;
        Hmat(i,i+1)=-t0x;
    elseif(ntc(i) == TRC)
        Hmat(i,i)=t0x+t0y;
        Hmat(i,i-n)=-t0y;
        Hmat(i,i-1)=-t0x;
    elseif(ntc(i) == MID)
        Hmat(i,i)=2*t0x+2*t0y;
        Hmat(i,i+1)=-t0x;
        Hmat(i,i-n)=-t0y;
        Hmat(i,i-1)=-t0x;
        Hmat(i,i+n)=-t0y;
    end
end;

%The below two lines are for the periodic potential case.
% Hmat(1,m)=-t0;
% Hmat(m,1)=-t0;

[alpha,E]=eig(Hmat);
E=diag(E);
[Enum,ind]=sort(E);

E1=E(ind(1));
Enum(1);
disp(E(1:5));
psi1=(alpha(:,ind(1)));
psi2=(alpha(:,ind(2)));
psi3=(alpha(:,ind(3)));
psi4=(alpha(:,ind(4)));
P1=psi1.*conj(psi1);
P2=psi2.*conj(psi2);
P3=psi3.*conj(psi3);
P4=psi4.*conj(psi4);

%electron density
n_density=P1+P2+P3+P4;


figure(1)
plot(x,P1,'green',x,P2,'blue',x,P3,'red',x,P4,'magenta');
title('Probability density for a given eigen state');
grid on;

figure(2)
plot(x,psi1,'green',x,psi2,'blue',x,psi3,'red',x,psi4,'magenta');
title('Eigen wave function');
grid on;